import numpy as np
import math
from scipy.spatial.transform import Rotation as R

# 设置Numpy的打印选项
# 精确位数3，不启用科学计数法
np.set_printoptions(precision=3, suppress=True)


a_points3d = np.loadtxt('ICP/a_points3d.txt', delimiter=',')
print(a_points3d)
b_points3d = np.loadtxt('ICP/b_points3d.txt', delimiter=',')
print(b_points3d)

def icp_solver_svd(a_points3d, b_points3d):
    '''计算的是从点集A到点集B的变换矩阵 ^{A}_{B}T
    P_A = ^{A}_{B}P * P_B
    '''
    # 分别计算两组点集的质心
    a_points3d_mean = np.mean(a_points3d, axis=0)
    b_points3d_mean = np.mean(b_points3d, axis=0)
    # 点集去除质心
    a_points3d_center = a_points3d - a_points3d_mean
    b_points3d_center = b_points3d - b_points3d_mean
    # 构造W矩阵
    W = np.zeros((3, 3), dtype="float32")
    for i in range(9):
        # V1 列向量
        v1 = a_points3d_center[i].reshape(-1, 1)
        # V2 行向量
        v2 = b_points3d_center[i].reshape(1, -1)
        # W矩阵累加
        W += v1.dot(v2)
    # SVD分解
    U,Sigma,Vt = np.linalg.svd(W)
    # 求解旋转矩阵
    R = U.dot(Vt)
    # 求解平移向量
    t = a_points3d_mean.reshape(-1,1) - R.dot(b_points3d_mean.reshape(-1, 1))
    # 构造变换矩阵
    # 工作台坐标系在Arm坐标系下的位姿描述
    return R, t
def rmat2euler(rmat):
    '''旋转矩阵转换为欧拉角'''
    alpha = None  # 偏航角
    beta = None  # 俯仰角
    gamma = None  # 横滚角

    r11, r12, r13, r21, r22, r23, r31, r32, r33 = rmat.reshape(-1)
    if abs(r31) >= (1 - 0.000001):
        # 出现万向锁的问题
        if r31 < 0:
            gamma = 0
            beta = np.pi / 2
            alpha = math.atan2(r23, r22)
            return [[gamma, beta, alpha]]
        else:
            gamma = 0
            beta = -np.pi / 2
            alpha = math.atan2(-r23, r22)
            return [[gamma, beta, alpha]]
    else:
        # 正常求解
        cos_beta = np.sqrt(r32 * r32 + r33 * r33)
        cos_beta_list = [cos_beta, -cos_beta]
        rpy_list = []
        for cos_beta in cos_beta_list:
            beta = math.atan2(-r31, cos_beta)
            #beta=math.degrees(format(beta,'f'))
            alpha = math.atan2(r21 / cos_beta, r11 / cos_beta)
            #alpha=math.degrees(format(alpha,'f'))
            gamma = math.atan2(r32 / cos_beta, r33 / cos_beta)
            #gamma=math.degrees(format(gamma,'f'))
            # beta = math.degrees(beta)
            # alpha = math.degrees(alpha)
            # gamma = math.degrees(gamma)
            rpy_list.append([gamma, beta, alpha])
        return rpy_list


# 求解旋转向量与平移向量
# 从点集
RH, t = icp_solver_svd(a_points3d, b_points3d)
print(f"旋转矩阵: \n{RH}")
print(f"平移向量: \n{t}")

for i in range(len(b_points3d)):
    # 获取点集A中的坐标
    pa = a_points3d[i].reshape((-1, 1))
    # 获取点集B中的坐标
    pb = b_points3d[i].reshape((-1, 1))
    # 空间变换
    pa2 = RH.dot(pb) + t
    print(f"[{i}]\n点集A中的点: {pa.reshape(-1)}  \n点集B中的点变换到A: {pa2.reshape(-1)} ")

R_laser_to_base = R.from_matrix(RH).as_euler('xyz', degrees=True)
print(f"平移向量: \n{t}")
print(R_laser_to_base)
